package data_structure.mooc.uf;

/**
 * 第三版并查集
 * 基于第二版并查树节点子节点的个数的优化
 */
public class UnionFind3 implements UnionFind{
    //parent数组记录的是元素的指向关系，从字节点指向父亲节点
    //有一个元素p，p指向q，则parent[p]=q;
    private int []parent;

    //sz[i]表示以i为根节点的树的节点的个数
    private int []sz;

    public UnionFind3(int size){
        parent = new int[size];
        sz=new int[size];
        for (int i = 0; i < size; i++) {
            parent[i]=i;
            sz[i]=1;
        }
    }
    @Override
    public int getSize() {
        return parent.length;
    }

    @Override
    public void unionElements(int p, int q) {
        int pRoot=find (p);
        int qRoot=find (q);

        if(pRoot==qRoot){
            return;
        }
        //p所在树的根节点指向q的根节点
        if(sz[pRoot]<sz[qRoot]){
            parent[pRoot]=qRoot;
            sz[qRoot]+=sz[pRoot];
        }else{
            parent[qRoot]=pRoot;
            sz[pRoot]+=sz[qRoot];
        }
    }

    /**
     * 返回元素p的根节点
     * @param p
     * @return
     */
    public int find(int p){
        if (p < 0 || p >= parent.length) {
            throw new IllegalArgumentException("Illegal index");
        }
        while(p!=parent[p]){
            p=parent[p];
        }
        return p;
    }
    @Override
    public boolean isConnected(int p, int q) {
        return find (p)==find (q);
    }
}
